from adapt.instance_based import TrAdaBoostR2
from adapt.instance_based import TwoStageTrAdaBoostR2
# import sys
from prettytable import PrettyTable
from sklearn import tree
from sklearn.model_selection import LeaveOneOut
from sklearn.model_selection import train_test_split as TTS
# from sko.AFSA import AFSA
from streamlit_extras.colored_header import colored_header
from streamlit_option_menu import option_menu

from business.algorithm.utils import *


def run():
    with st.sidebar:
        sub_option = option_menu(None, ["提升方法"])
    if sub_option == "提升方法":
        colored_header(label="样本迁移学习", description=" ", color_name="violet-90")

        file = st.file_uploader("Upload `.csv`file", type=['csv'], label_visibility="collapsed",
                                accept_multiple_files=True)
        if len(file) < 2:
            table = PrettyTable(['file name', 'class', 'descirption'])
            table.add_row(['file_1', 'target_data', 'target domain'])
            table.add_row(['file_2', 'source_data_1', '1 source domain'])
            table.add_row(['...', '...', '...'])
            table.add_row(['file_n', 'source_data_n', 'n source domain'])
            st.write(table)

        elif len(file) >= 2:
            df_target = pd.read_csv(file[0])
            df_source = pd.read_csv(file[1])

            colored_header(label="数据信息", description=" ", color_name="violet-70")

            nrow = st.slider("rows", 1, len(df_target), 5)
            df_nrow = df_target.head(nrow)
            st.write(df_nrow)
            colored_header(label="特征&目标", description=" ", color_name="violet-70")

            target_num = st.number_input('目标数量', min_value=1, max_value=10, value=1)

            col_feature, col_target = st.columns(2)
            # features
            target_features = df_target.iloc[:, :-target_num]
            source_features = df_source.iloc[:, :-target_num]
            # targets
            target_targets = df_target.iloc[:, -target_num:]
            source_targets = df_source.iloc[:, -target_num:]

            with col_feature:
                st.write(target_features.head())
            with col_target:
                st.write(target_targets.head())
            # =================== model ====================================
            reg = REGRESSOR(target_features, target_targets)

            reg.td_features = target_features
            reg.td_targets = target_targets
            reg.sd_features = source_features
            reg.sd_targets = source_targets

            colored_header(label="target", description=" ", color_name="violet-70")

            target_selected_option = st.selectbox('target', list(reg.targets)[::-1])

            reg.td_targets = target_targets[target_selected_option]
            reg.sd_targets = source_targets[target_selected_option]

            colored_header(label="Transfer", description=" ", color_name="violet-30")

            model_path = './models/transfer learning'

            template_alg = model_platform(model_path)

            inputs, col2 = template_alg.show()
            if inputs['model'] == 'TrAdaboostR2':
                DTR = tree.DecisionTreeRegressor()
                with col2:
                    with st.expander('Operator'):
                        operator = st.selectbox('', ('train test split', 'cross val score', 'leave one out'),
                                                label_visibility="collapsed")
                        if operator == 'train test split':
                            inputs['test size'] = st.slider('test size', 0.1, 0.5, 0.2)
                            reg.Xtrain, reg.Xtest, reg.Ytrain, reg.Ytest = TTS(reg.td_features, reg.td_targets,
                                                                               test_size=inputs['test size'],
                                                                               random_state=inputs['random state'])

                        elif operator == 'cross val score':
                            cv = st.number_input('cv', 1, 20, 5)

                        elif operator == 'leave one out':
                            loo = LeaveOneOut()

                colored_header(label="Training", description=" ", color_name="violet-30")
                with st.container():
                    button_train = st.button('Train', use_container_width=True)
                if button_train:
                    if operator == 'train test split':
                        reg.model = TrAdaBoostR2(DTR, n_estimators=inputs['n_estimators'], Xt=reg.Xtrain, yt=reg.Ytrain,
                                                 verbose=-1)

                        reg.TrAdaBoostR2()

                        result_data = pd.concat([reg.Ytest, pd.DataFrame(reg.Ypred)], axis=1)
                        result_data.columns = ['actual', 'prediction']
                        plot_and_export_results(reg, "TrAdaboostR2")

                    elif operator == 'cross val score':

                        kf = KFold(n_splits=5, shuffle=True, random_state=42)

                        y_pred_list = []
                        y_test_list = []

                        for train_index, test_index in kf.split(reg.td_features.values):
                            X_train, X_test = reg.td_features.values[train_index], reg.td_features.values[test_index]
                            y_train, y_test = reg.td_targets.values[train_index], reg.td_targets.values[test_index]
                            y_train = y_train.reshape(-1, 1)
                            y_test = y_test.reshape(-1, 1)

                            # 创建模型并训练
                            model = TrAdaBoostR2(DTR, n_estimators=inputs['n_estimators'], Xt=X_train, yt=y_train,
                                                 verbose=-1)
                            model.fit(reg.sd_features.values, reg.sd_targets.values.reshape(-1, 1))

                            y_pred = model.predict(X_test)
                            y_pred_list.append(y_pred)
                            y_test_list.append(y_test)

                        Y_pred = np.concatenate(y_pred_list)
                        Y_test = np.concatenate(y_test_list)
                        st.write(f'R2: {r2_score(Y_pred, Y_test)}')
                        fig, ax = plt.subplots(figsize=(5, 4))
                        ax.scatter(Y_pred, Y_test, marker='o', color='#000080', zorder=1, facecolors='none')
                        lims = [
                            np.min([ax.get_xlim(), ax.get_ylim()]),  # min of both axes
                            np.max([ax.get_xlim(), ax.get_ylim()]),  # max of both axes
                        ]
                        ax.tick_params(direction='in', length=5)
                        ax.plot(lims, lims, zorder=8, linewidth=2, linestyle='solid', color='#FF0000')
                        ax.set_xlim(lims)
                        ax.set_ylim(lims)
                        plt.xlabel("Actual")
                        plt.ylabel("Prediction")
                        st.pyplot(fig)

                        with st.expander("model"):
                            tmp_download_link = download_button(model, 'TaAdaboostR2_cv' + '.pickle',
                                                                button_text='download')
                            st.markdown(tmp_download_link, unsafe_allow_html=True)
                        result_data = pd.concat([pd.DataFrame(Y_test), pd.DataFrame(Y_pred)], axis=1)
                        result_data.columns = ['actual', 'prediction']
                        with st.expander('prediction'):
                            st.write(result_data)
                            tmp_download_link = download_button(result_data, f'prediction.csv', button_text='download')
                            st.markdown(tmp_download_link, unsafe_allow_html=True)

                    elif operator == 'leave one out':

                        y_pred_list = []
                        y_test_list = []
                        for train_index, test_index in loo.split(reg.td_features.values):
                            X_train, X_test = reg.td_features.values[train_index], reg.td_features.values[test_index]
                            y_train, y_test = reg.td_targets.values[train_index], reg.td_targets.values[test_index]
                            y_train = y_train.reshape(-1, 1)
                            y_test = y_test.reshape(-1, 1)
                            # 创建模型并训练
                            model = TrAdaBoostR2(DTR, n_estimators=inputs['n_estimators'], Xt=X_train, yt=y_train,
                                                 verbose=-1)
                            model.fit(reg.sd_features.values, reg.sd_targets.values.reshape(-1, 1))
                            y_pred = model.predict(X_test)
                            y_pred_list.append(y_pred)
                            y_test_list.append(y_test)

                        Y_pred = np.ravel(y_pred_list)
                        Y_test = np.ravel(y_test_list)
                        st.write(f'R2: {r2_score(Y_pred, Y_test)}')
                        fig, ax = plt.subplots(figsize=(5, 4))
                        ax.scatter(Y_pred, Y_test, marker='o', color='#000080', zorder=1, facecolors='none')
                        lims = [
                            np.min([ax.get_xlim(), ax.get_ylim()]),  # min of both axes
                            np.max([ax.get_xlim(), ax.get_ylim()]),  # max of both axes
                        ]
                        ax.tick_params(direction='in', length=5)
                        ax.plot(lims, lims, zorder=8, linewidth=2, linestyle='solid', color='#FF0000')
                        ax.set_xlim(lims)
                        ax.set_ylim(lims)
                        plt.xlabel("Actual")
                        plt.ylabel("Prediction")
                        st.pyplot(fig)

                        with st.expander("model"):
                            tmp_download_link = download_button(model, 'TaAdaboostR2_loo' + '.pickle',
                                                                button_text='download')
                            st.markdown(tmp_download_link, unsafe_allow_html=True)
                        result_data = pd.concat([pd.DataFrame(Y_test), pd.DataFrame(Y_pred)], axis=1)
                        result_data.columns = ['actual', 'prediction']
                        with st.expander('prediction'):
                            st.write(result_data)
                            tmp_download_link = download_button(result_data, f'prediction.csv', button_text='download')
                            st.markdown(tmp_download_link, unsafe_allow_html=True)


            elif inputs['model'] == 'TwoStageTrAdaboostR2':
                DTR = tree.DecisionTreeRegressor()
                with col2:
                    with st.expander('Operator'):
                        operator = st.selectbox('', ('train test split', 'cross val score', 'leave one out'),
                                                label_visibility="collapsed")
                        if operator == 'train test split':
                            inputs['test size'] = st.slider('test size', 0.1, 0.5, 0.2)
                            reg.Xtrain, reg.Xtest, reg.Ytrain, reg.Ytest = TTS(reg.td_features, reg.td_targets,
                                                                               test_size=inputs['test size'],
                                                                               random_state=inputs['random state'])

                        elif operator == 'cross val score':
                            cv = st.number_input('cv', 1, 20, 5)

                        elif operator == 'leave one out':
                            loo = LeaveOneOut()

                colored_header(label="Training", description=" ", color_name="violet-30")
                with st.container():
                    button_train = st.button('Train', use_container_width=True)
                if button_train:
                    if operator == 'train test split':
                        reg.model = TwoStageTrAdaBoostR2(DTR, n_estimators=inputs['n_estimators'], Xt=reg.Xtrain,
                                                         yt=reg.Ytrain, verbose=-1)

                        reg.TwoStageTrAdaBoostR2()

                        result_data = pd.concat([reg.Ytest, pd.DataFrame(reg.Ypred)], axis=1)
                        result_data.columns = ['actual', 'prediction']
                        plot_and_export_results(reg, "TrAdaboostR2")

                    elif operator == 'cross val score':

                        kf = KFold(n_splits=5, shuffle=True, random_state=42)

                        y_pred_list = []
                        y_test_list = []

                        for train_index, test_index in kf.split(reg.td_features.values):
                            X_train, X_test = reg.td_features.values[train_index], reg.td_features.values[test_index]
                            y_train, y_test = reg.td_targets.values[train_index], reg.td_targets.values[test_index]
                            y_train = y_train.reshape(-1, 1)
                            y_test = y_test.reshape(-1, 1)

                            # 创建模型并训练
                            model = TwoStageTrAdaBoostR2(DTR, n_estimators=inputs['n_estimators'], Xt=X_train,
                                                         yt=y_train, verbose=-1)
                            model.fit(reg.sd_features.values, reg.sd_targets.values.reshape(-1, 1))

                            y_pred = model.predict(X_test)
                            y_pred_list.append(y_pred)
                            y_test_list.append(y_test)

                        Y_pred = np.concatenate(y_pred_list)
                        Y_test = np.concatenate(y_test_list)
                        st.write(f'R2: {r2_score(Y_pred, Y_test)}')
                        fig, ax = plt.subplots(figsize=(5, 4))
                        ax.scatter(Y_pred, Y_test, marker='o', color='#000080', zorder=1, facecolors='none')
                        lims = [
                            np.min([ax.get_xlim(), ax.get_ylim()]),  # min of both axes
                            np.max([ax.get_xlim(), ax.get_ylim()]),  # max of both axes
                        ]
                        ax.tick_params(direction='in', length=5)
                        ax.plot(lims, lims, zorder=8, linewidth=2, linestyle='solid', color='#FF0000')
                        ax.set_xlim(lims)
                        ax.set_ylim(lims)
                        plt.xlabel("Actual")
                        plt.ylabel("Prediction")
                        st.pyplot(fig)

                        with st.expander("model"):
                            tmp_download_link = download_button(model, 'TaAdaboostR2_cv' + '.pickle',
                                                                button_text='download')
                            st.markdown(tmp_download_link, unsafe_allow_html=True)
                        result_data = pd.concat([pd.DataFrame(Y_test), pd.DataFrame(Y_pred)], axis=1)
                        result_data.columns = ['actual', 'prediction']
                        with st.expander('prediction'):
                            st.write(result_data)
                            tmp_download_link = download_button(result_data, f'prediction.csv', button_text='download')
                            st.markdown(tmp_download_link, unsafe_allow_html=True)

                    elif operator == 'leave one out':

                        y_pred_list = []
                        y_test_list = []
                        for train_index, test_index in loo.split(reg.td_features.values):
                            X_train, X_test = reg.td_features.values[train_index], reg.td_features.values[test_index]
                            y_train, y_test = reg.td_targets.values[train_index], reg.td_targets.values[test_index]
                            y_train = y_train.reshape(-1, 1)
                            y_test = y_test.reshape(-1, 1)
                            # 创建模型并训练
                            model = TwoStageTrAdaBoostR2(DTR, n_estimators=inputs['n_estimators'], Xt=X_train,
                                                         yt=y_train, verbose=-1)
                            model.fit(reg.sd_features.values, reg.sd_targets.values.reshape(-1, 1))
                            y_pred = model.predict(X_test)
                            y_pred_list.append(y_pred)
                            y_test_list.append(y_test)

                        Y_pred = np.ravel(y_pred_list)
                        Y_test = np.ravel(y_test_list)
                        st.write(f'R2: {r2_score(Y_pred, Y_test)}')
                        fig, ax = plt.subplots(figsize=(5, 4))
                        ax.scatter(Y_pred, Y_test, marker='o', color='#000080', zorder=1, facecolors='none')
                        lims = [
                            np.min([ax.get_xlim(), ax.get_ylim()]),  # min of both axes
                            np.max([ax.get_xlim(), ax.get_ylim()]),  # max of both axes
                        ]
                        ax.tick_params(direction='in', length=5)
                        ax.plot(lims, lims, zorder=8, linewidth=2, linestyle='solid', color='#FF0000')
                        ax.set_xlim(lims)
                        ax.set_ylim(lims)
                        plt.xlabel("Actual")
                        plt.ylabel("Prediction")
                        st.pyplot(fig)

                        with st.expander("model"):
                            tmp_download_link = download_button(model, 'TaAdaboostR2_loo' + '.pickle',
                                                                button_text='download')
                            st.markdown(tmp_download_link, unsafe_allow_html=True)
                        result_data = pd.concat([pd.DataFrame(Y_test), pd.DataFrame(Y_pred)], axis=1)
                        result_data.columns = ['actual', 'prediction']
                        with st.expander('prediction'):
                            st.write(result_data)
                            tmp_download_link = download_button(result_data, f'prediction.csv', button_text='download')
                            st.markdown(tmp_download_link, unsafe_allow_html=True)
                st.write('---')